from polyhedron import Polyhedron
from carbon import Carbon
import pandas as pd
import numpy as np

#单点能
single_energy = {'B':-24.579564, 'C':-37.73266, 'N':-54.527285}

def read_excel_energy(file_name):
    data = pd.read_excel(file_name, sheet_name='98')
    return data['能量'].to_numpy()


def prepare_carbon_data():
    carbon_list = Carbon.from_bin_all(r'D:\360安全云盘同步版\Cluster\C\C60-160\full60_160\Full_codes_98_ipr')
    m,n = len(carbon_list), 98
    d = np.empty((m,12,12), dtype='int')
    carbon_list = [p.dual()[0] for p in carbon_list]
    for i,c in enumerate(carbon_list):
        idx5 = list(filter(lambda v:c.degree(v)==5, range(len(c.at))))
        dist = c.dist()
        d[i] = dist[idx5][:,idx5]
        print(i)
    e = read_excel_energy(r'D:\360安全云盘同步版\Cluster\C\C60-160\C60_116-energy.xlsx')
    e = (single_energy['C'] -e/n) * 27.2114
    np.savez('C98_DE.npz',D=d,E=e)
    
    
def prepare_BN_data(folder_name, size):
    '''将文件夹下所有子文件夹中的outmol文件中的结构和能量读出
    保存能量、三角多面体的邻接矩阵、顶点哈希3个数组'''
    Atom.tolerance_max = 1.25
    
    def prepare_data_one(file_name):
        c = read(file_name)
        p = Polyhedron.from_cluster(c)
        assert set([len(t) for t in p.at]) == set([3])
        q,_ = p.dual() #转为三角多面体
        c.set_energy(Dmol3.read_energy(file_name))
        e = c.binding_energy(single_energy)
        return q,e #能量，邻接表矩阵
       
    E = []
    tri_poly = []
    fp = open('BN%d.txt'%size,'w')
    ff = open('BN%d_list.txt'%size,'w')
    n = 0
    for folder in os.listdir(folder_name):
        path = os.path.join(folder_name,folder)
        if os.path.isfile(path):
            continue
        for file_name in os.listdir(path):
        #for f in glob(folder_name+'/*.outmol'):
            if not file_name.endswith('.outmol'):
                continue
            f = os.path.join(path,file_name)
            print(f)
            q,e = prepare_data_one(f)
            assert sum(set([len(t) for t in q.at])) in [10,14,18] #4,6,8
            tri_poly.append(q)
            E.append(e)
            fp.write(str(q))
            ff.write(f+'\n')
            n += 1
     
    print(len(tri_poly))
    E = np.array(E, dtype='float32')
    size = len(tri_poly[0].at)
    H = np.empty((n,size), dtype='int') #顶点哈希
    A = np.empty((n,size,size), dtype='int')
    for i,q in enumerate(tri_poly):
        H[i] = np.array([HASH_DICT[q._vertex_hash(v)] for v in range(len(q.at))])
        A[i] = q.to_adj()
        
    np.savez('BN%d_aei.npz'%size,A=A,E=E,I=H)


#read_excel_energy(r'D:\360安全云盘同步版\Cluster\C\C60-160\C60_116-energy.xlsx')
prepare_BN_data()